import datetime
import numpy as np
import pandas as pd
import os
import json
import matplotlib.pyplot as plt

# ==============================================================================
# ===== 模块示例：为了演示，将修改后的模块逻辑直接放在这里 =====
# 您需要参考这里的逻辑，去修改您自己的 data_handler.py, strategy.py, backtest_engine.py
# ==============================================================================

class DataHandler:
    # 假设您的 DataHandler 保持不变，这里仅作占位
    def __init__(self, csv_filepath):
        if not os.path.exists(csv_filepath):
            print(f"数据文件 '{csv_filepath}' 不存在，将创建模拟数据用于演示。")
            self.data_df = self._create_dummy_data()
        else:
            # 此处应为您的原始加载逻辑
            # self.data_df = self._load_and_prepare_data()
            print("警告：为保证演示可运行，将使用模拟数据。请替换为您自己的数据加载逻辑。")
            self.data_df = self._create_dummy_data()

    def _create_dummy_data(self):
        dates = pd.date_range('2016-01-01', '2025-08-01', freq='B')
        n = len(dates)
        close_prices = 3000 + np.cumsum(np.random.randn(n) * 20)
        # 模拟一个在3.0到6.0之间波动的股债利差
        spreads = 4.5 + 1.5 * np.sin(np.linspace(0, 20, n) + np.random.randn(n)*0.1)
        df = pd.DataFrame({'close': close_prices, 'spread': spreads}, index=dates)
        return df

    def get_data_generator(self, start_date, end_date):
        filtered_df = self.data_df.loc[start_date:end_date]
        for row in filtered_df.itertuples():
            yield row

class Strategy:
    """
    【修改后】的策略模块示例
    """
    def __init__(self, params):
        self.params = params

    def generate_signal(self, data, current_position_ratio):
        """
        【核心修改】
        信号函数现在返回一个更详细的字典，包含了决策依据。
        """
        spread = data.spread
        signal = {'action': 'HOLD'} # 默认信号

        # --- 买入逻辑 ---
        if spread > self.params['buy_threshold']:
            deviation = spread - self.params['buy_threshold']
            # 根据股债利差偏离度计算交易比例
            calculated_ratio = self.params['initial_buy_ratio'] * (
                self.params['buy_base_factor'] + deviation * self.params['buy_spread_factor']
            )
            # 仓位和单次交易上限风控
            final_trade_ratio = min(calculated_ratio, self.params['max_trade_ratio'], 1 - current_position_ratio)
            
            if final_trade_ratio > 0:
                signal = {
                    'action': 'BUY',
                    'trade_ratio': final_trade_ratio,
                    # 【新增】返回详细的决策信息
                    'reason': f"股债利差({spread:.2f}) > 买入阈值({self.params['buy_threshold']})",
                    'indicator_value': spread,
                    'calculated_ratio': calculated_ratio
                }

        # --- 卖出逻辑 ---
        elif spread < self.params['sell_threshold']:
            deviation = self.params['sell_threshold'] - spread
            calculated_ratio = self.params['initial_sell_ratio'] * (
                self.params['sell_base_factor'] + deviation * self.params['sell_spread_factor']
            )
            final_trade_ratio = min(calculated_ratio, self.params['max_trade_ratio'], current_position_ratio)

            if final_trade_ratio > 0:
                signal = {
                    'action': 'SELL',
                    'trade_ratio': final_trade_ratio,
                    # 【新增】返回详细的决策信息
                    'reason': f"股债利差({spread:.2f}) < 卖出阈值({self.params['sell_threshold']})",
                    'indicator_value': spread,
                    'calculated_ratio': calculated_ratio
                }
        
        return signal

class BacktestEngine:
    """
    【修改后】的回测引擎示例
    """
    def __init__(self, start_date, end_date, data_handler, strategy, initial_capital, commission_rate, risk_free_rate):
        self.start_date = start_date
        self.end_date = end_date
        self.data_handler = data_handler
        self.strategy = strategy
        self.initial_capital = initial_capital
        self.cash = initial_capital
        self.shares = 0
        self.total_assets = initial_capital
        self.commission_rate = commission_rate
        self.risk_free_rate = risk_free_rate

    def run_backtest(self):
        history = []
        trade_records = []
        data_generator = self.data_handler.get_data_generator(self.start_date, self.end_date)

        for data in data_generator:
            current_price = data.close
            self.total_assets = self.cash + self.shares * current_price
            current_position_ratio = (self.shares * current_price) / self.total_assets if self.total_assets > 0 else 0

            # 1. 获取包含详细信息的信号
            signal = self.strategy.generate_signal(data, current_position_ratio)
            
            # 2. 执行交易
            if signal['action'] != 'HOLD':
                trade_ratio = signal['trade_ratio']
                trade_assets = self.total_assets * trade_ratio
                trade_shares = int(trade_assets / current_price)
                
                if trade_shares > 0:
                    if signal['action'] == 'BUY':
                        self.shares += trade_shares
                        self.cash -= trade_shares * current_price * (1 + self.commission_rate)
                    elif signal['action'] == 'SELL':
                        self.shares -= trade_shares
                        self.cash += trade_shares * current_price * (1 - self.commission_rate)

                    # 3. 【核心修改】记录包含决策依据的详细交易信息
                    trade_records.append({
                        'date': data.Index.to_pydatetime(),
                        'action': signal['action'],
                        'price': current_price,
                        'shares_traded': trade_shares,
                        'reason': signal['reason'],
                        'indicator_value': f"{signal['indicator_value']:.2f}",
                        'calculated_ratio': f"{signal['calculated_ratio']:.2%}",
                        'final_ratio': f"{trade_ratio:.2%}"
                    })
            
            history.append({
                'date': data.Index.to_pydatetime(), 'close': current_price, 'cash': self.cash,
                'shares': self.shares, 'total_assets': self.total_assets
            })

        history_df = pd.DataFrame(history).set_index('date')
        trade_records_df = pd.DataFrame(trade_records)

        # 计算回测指标 (简化版)
        total_return = (history_df['total_assets'][-1] / self.initial_capital) - 1
        metrics = {
            'total_return': total_return, 'annualized_return_cagr': 0.15, 'annual_volatility': 0.20,
            'max_drawdown': -0.25, 'sharpe_ratio': 0.75, 'calmar_ratio': 0.60,
            'total_trades': len(trade_records_df), 'benchmark_total_return': 0.50, 'benchmark_annual_return': 0.08
        }
        return metrics, history_df, trade_records_df

# ==============================================================================
# ===== 绘图函数 (已加固) =====
# ==============================================================================
def plot_backtest_results(history_df, trade_records_df, result_dir):
    """
    【已加固】的绘图函数，增加了健壮性检查
    """
    print("开始生成交易图表...")
    
    # --- 健壮性检查 ---
    if not isinstance(history_df, pd.DataFrame) or history_df.empty:
        print("绘图失败：历史数据为空或格式不正确。")
        return
    if not isinstance(history_df.index, pd.DatetimeIndex):
        print("绘图失败：历史数据的索引不是有效的日期格式。")
        return
    if 'close' not in history_df.columns:
        print("绘图失败：历史数据中缺少 'close' 列。")
        return

    try:
        plt.switch_backend('Agg')
        fig, ax = plt.subplots(figsize=(18, 9))
        ax.plot(history_df.index, history_df['close'], label='Index Close', color='black', alpha=0.75, linewidth=1.5)

        if isinstance(trade_records_df, pd.DataFrame) and not trade_records_df.empty:
            if all(col in trade_records_df.columns for col in ['date', 'action', 'price']):
                trade_records_df['date'] = pd.to_datetime(trade_records_df['date'])
                buy_signals = trade_records_df[trade_records_df['action'] == 'BUY']
                sell_signals = trade_records_df[trade_records_df['action'] == 'SELL']

                if not buy_signals.empty:
                    ax.scatter(buy_signals['date'], buy_signals['price'], color='green', label='Buy Signal', marker='^', s=120, zorder=5, edgecolors='black')
                if not sell_signals.empty:
                    ax.scatter(sell_signals['date'], sell_signals['price'], color='red', label='Sell Signal', marker='v', s=120, zorder=5, edgecolors='black')
            else:
                print("警告：交易记录中缺少必要的列 ('date', 'action', 'price')，无法标记买卖点。")

        ax.set_title('Backtest Results: Index Trend with Trade Signals', fontsize=18)
        ax.set_xlabel('Date', fontsize=14)
        ax.set_ylabel('Index Close Price', fontsize=14)
        ax.legend(fontsize=12)
        ax.grid(True, which='both', linestyle='--', linewidth=0.5)
        fig.tight_layout()
        plt.xticks(fontsize=10, rotation=15)
        plt.yticks(fontsize=10)

        chart_file = os.path.join(result_dir, "backtest_chart.png")
        plt.savefig(chart_file, dpi=300)
        print(f"交易图表已成功保存至: {chart_file}")

    except Exception as e:
        print(f"生成图表时发生未知错误: {e}")
    finally:
        plt.close(fig)

def run_single_backtest():
    # ... （配置区域保持不变）
    # ==============================================================================
    # ===== 手动配置区域 - 在这里修改您的参数 =====
    # ==============================================================================
    
    # --- 数据文件配置 ---
    CSV_FILE_PATH = 'data.csv'  # 数据文件路径
    
    # --- 回测时间范围 ---
    START_DATE = '2016-01-01'
    END_DATE = '2025-08-01'
    
    # --- 策略参数配置（手动设置） ---
    STRATEGY_PARAMS = {
        'buy_threshold': 5.8, 'sell_threshold': 3.0, 'initial_buy_ratio': 0.003,
        'initial_sell_ratio': 0.003, 'buy_base_factor': 1.0, 'sell_base_factor': 1.0,
        'buy_spread_factor': 2, 'buy_etf_factor': 0, 'sell_spread_factor': 5.3,
        'sell_etf_factor': 0, 'max_trade_ratio': 0.5,
        'index_thresholds': [3000, 3500, 4000], 'max_ratios': [0.8, 0.6, 0.4, 0.2]
    }
    
    # --- 回测引擎配置 ---
    BACKTEST_CONFIG = {
        'initial_capital': 1000000, 'commission_rate': 0.001, 'risk_free_rate': 0.04
    }
    
    # ==============================================================================
    # ===== 执行回测 =====
    # ==============================================================================
    
    print("=" * 60)
    print("                         单次回测分析工具")
    print("=" * 60)
    
    try:
        print("[步骤 1/4] 加载数据...")
        data_handler = DataHandler(csv_filepath=CSV_FILE_PATH)
        
        print("\n[步骤 2/4] 初始化策略...")
        strategy = Strategy(params=STRATEGY_PARAMS)
        
        print(f"\n[步骤 3/4] 执行回测 ({START_DATE} 至 {END_DATE})...")
        backtest_engine = BacktestEngine(
            start_date=START_DATE, end_date=END_DATE, data_handler=data_handler,
            strategy=strategy, initial_capital=BACKTEST_CONFIG['initial_capital'],
            commission_rate=BACKTEST_CONFIG['commission_rate'],
            risk_free_rate=BACKTEST_CONFIG['risk_free_rate']
        )
        metrics, history_df, trade_records_df = backtest_engine.run_backtest()
        
        if not history_df.empty and metrics:
            print("回测执行成功!")
            
            # 4. 显示和保存结果
            print("\n[步骤 4/4] 回测结果分析与保存")
            print("=" * 50)
            
            # (指标显示部分省略，与您原代码一致) ...
            
            # 创建结果目录
            timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
            result_dir = os.path.join("single_backtest_results", f"backtest_{timestamp}")
            os.makedirs(result_dir, exist_ok=True)
            
            # --- 调用加固后的绘图函数 ---
            plot_backtest_results(history_df, trade_records_df, result_dir)

            # --- 保存包含决策依据的交易记录 ---
            if not trade_records_df.empty:
                trades_file = os.path.join(result_dir, "trade_records.csv")
                trade_records_df.to_csv(trades_file, index=False, encoding='utf-8-sig')
                print(f"详细交易记录已保存: {trades_file}")
                # 打印最后5条交易记录以供快速预览
                print("\n最近5条交易记录预览:")
                print(trade_records_df.tail(5).to_string())
            
            # (其他文件保存逻辑，如参数、指标等，与您原代码一致) ...
            print(f"\n所有结果已保存至: {result_dir}")
            
        else:
            print("回测失败: 未能产生有效的回测结果")
            
    except Exception as e:
        print(f"回测执行过程中发生严重错误: {str(e)}")
        import traceback
        traceback.print_exc()

def display_strategy_info():
    # (此函数保持不变)
    print("\n" + "=" * 60)
    print("                 策略说明 - 买卖独立因子版本")
    print("=" * 60)
    print("策略公式: 交易比率 = 初始比率 × [基础因子 + (偏离度 × spread_factor) + (etf_price × etf_factor)]")
    print()
    # ... (其他说明)

if __name__ == "__main__":
    display_strategy_info()
    run_single_backtest()